#include "framework.hh"
#include "framework.hh"
#include "framework.hh"
#include "module_frame.hh"
#include "entity.hh"
#include "entity_search_cache.hh"
#include "flatbuffers_helper.hh"
#include "../protocol/protocol_header.hh"
#include "../protocol/protocol_register.hh"
#include "../../corelib/logger/log.hh"
#include "../protocol/locator/register_entity_req_generated.h"
#include "../protocol/locator/register_entity_ack_generated.h"
#include "../protocol/locator/unregister_entity_req_generated.h"
#include "../protocol/locator/entity_relay_msg_req_generated.h"
#include "../protocol/locator/entity_search_req_generated.h"
#include "../../corelib/grid/thirdparty/flatbuffers/flatbuffers.h"

void Framework::onEntityData(std::uint64_t from, std::uint64_t to, std::uint32_t msgID, Message & msg) {
    auto handler = ProtocolRegisterRef.getHandler(msgID);
    if (handler) {
        handler->onEntityMessage(from, to, msg);
    } else {
        errorlog << "Message ID not found, ENTITY MSGID:" << msgID << endlog;
    }
}

void Framework::onLocatorRehash(PathID pathID) {
    static EntityRegistryVector entityRegistries;
    entityRegistries.clear();
    for (auto it = registeredEntityMap_.begin(); it != registeredEntityMap_.end();) {
        if (it->second.pathID == pathID) {
            // Remove the entity belongs to this locator
            entityRegistries.push_back(it->second);
            registeredEntityMap_.erase(it++);
        } else {
            it++;
        }
    }
    warnlog << "Rehash locator, path ID:" << pathID << ", rehash count:" << entityRegistries.size() << endlog;
    for (auto registry : entityRegistries) {
        // register to other locator
        registerEntity(registry.entityID, registry.entityType);
    }
}

bool Framework::talkLocator(PathID pathID, std::uint32_t msgID, flatbuffers::FlatBufferBuilder & builder,
    std::uint32_t userMsgID, const Message & msg) {
    auto bytes = msg.ByteSize();
    auto headerOffset = sizeof(ProtocolLocatorHeader) + static_cast<std::uint32_t>(builder.GetSize());
    auto msgBodyOffset = sizeof(ProtocolLocatorHeader) + builder.GetSize() + sizeof(ProtocolHeader);
    auto flatMsgOffset = sizeof(ProtocolLocatorHeader);
    auto totalBytes = msgBodyOffset + bytes;
    FRAMEWORK_CHECK_RETURN(totalBytes > kratos::NETWORK_IOLOOP_RECV_BUFFER_MAX_SIZE, false,
        "Packet size invalid, msgID:" << msgID << ", size:" << bytes);
    buildLocatorHeaderInSendBuffer(0, msgID, builder.GetSize(), totalBytes);
    buildHeaderInSendBuffer(headerOffset, userMsgID, sizeof(ProtocolHeader) + bytes);
    memcpy(sendBuffer_.get() + flatMsgOffset, reinterpret_cast<char*>(builder.GetBufferPointer()),
        builder.GetSize());
    FRAMEWORK_CHECK_RETURN(!msg.SerializeToArray(sendBuffer_.get() + msgBodyOffset, bytes), false,
        "SerializeToArray failed, msgID:" << msgID << ", size:" << bytes);
    return talkPath(pathID, sendBuffer_.get(), totalBytes);
}

bool Framework::tryToStartRegisterTimer(std::uint64_t entityID) {
    // Start timer for unfinished registering
    auto timerHandle = startTimer(1000, [&](TimerHandle timerID, UserData userData) {
        auto pendingEntityID = static_cast<std::uint64_t>(userData);
        auto it = pendingEntityMap_.find(pendingEntityID);
        if (it == pendingEntityMap_.end()) {
            cancelTimer(timerID);
            return;
        } else {
            if (it->second.retryTimes > MAX_RETRY_TIMES) {
                // Max retry times
                errorlog << "Entity registering reaches max retry times, entity ID:" << pendingEntityID << endlog;
                pendingEntityMap_.erase(it);
                cancelTimer(timerID);
                return;
            } else {
                it->second.retryTimes += 1;
            }
            auto pathID = ModuleRef.getEngine()->getLocator(pendingEntityID);
            if (pathID != INVALID_PATH_ID) {
                BUILD_FLATBUFFER(RegisterEntityReq, fb, pendingEntityID, it->second.entityType);
                talkLocator(pathID, getMsgRealID(LocatorMsgID::REGISTER_REQ), fb);
            }
        }
    }, entityID);
    FRAMEWORK_CHECK_RETURN(timerHandle == INVALID_TIMER_ID, false, "Start registering timer failed, entity ID:" << entityID);
    return true;
}

void Framework::onFinishRegister(PathID pathID, std::uint64_t entityID) {
    auto it = pendingEntityMap_.find(entityID);
    if (it == pendingEntityMap_.end()) {
        return;
    }
    //verblog << "Entity registered, entity ID:" << entityID << endlog;
    it->second.pathID = pathID;
    registeredEntityMap_.emplace(it->first, it->second);
    pendingEntityMap_.erase(it);
}

ProtocolHeader * Framework::buildHeaderInSendBuffer(std::size_t offset, std::uint32_t msgID, std::uint32_t length) {
    auto header = reinterpret_cast<ProtocolHeader*>(sendBuffer_.get() + offset);
    header->length = length;
    header->msgID = msgID;
    header->hton();
    return header;
}

ProtocolLocatorHeader * Framework::buildLocatorHeaderInSendBuffer(std::size_t offset, std::uint32_t msgID,
    std::uint32_t locatorMsgLength, std::uint32_t length) {
    auto locatorHeader  = reinterpret_cast<ProtocolLocatorHeader*>(sendBuffer_.get() + offset);
    locatorHeader->msgID = msgID;
    locatorHeader->locatorMsgLength = locatorMsgLength;
    locatorHeader->length = length;
    locatorHeader->hton();
    return locatorHeader;
}

bool Framework::talkLocator(PathID pathID, std::uint32_t msgID, flatbuffers::FlatBufferBuilder & builder) {
    buildLocatorHeaderInSendBuffer(0, msgID, builder.GetSize(),
        sizeof(ProtocolLocatorHeader) + static_cast<std::uint32_t>(builder.GetSize()));
    memcpy(sendBuffer_.get() + sizeof(ProtocolLocatorHeader), reinterpret_cast<char*>(builder.GetBufferPointer()),
        builder.GetSize());
    return talkPath(pathID, sendBuffer_.get(), sizeof(ProtocolLocatorHeader) + builder.GetSize());
}
